import torch
import numpy as np
import os
from model import EnhancedWorldModel
from utils import generate_dataset
from train import train, prepare_data_loaders
from evaluate import evaluate, run_test_episodes
from config import *


def main():
    # 设置随机种子
    torch.manual_seed(RANDOM_SEED)
    np.random.seed(RANDOM_SEED)

    # 创建模型目录
    os.makedirs("models", exist_ok=True)

    # 计算额外特征带来的输入尺寸增加
    additional_features = 0
    if USE_DISTANCE_MAPS:
        additional_features += 3  # 障碍物、目标、智能体距离图
    if USE_DIRECTION_MAPS:
        additional_features += 4  # 目标方向(sin,cos)、智能体方向(sin,cos)

    total_node_features = CELL_FEATURES + 2 + additional_features  # 基本特征 + 位置编码 + 额外特征

    # 生成数据集
    print("Generating dataset...")
    print(f"使用增强状态表示: 距离图={USE_DISTANCE_MAPS}, 方向图={USE_DIRECTION_MAPS}")
    train_data, val_data, test_data = generate_dataset(num_samples=NUM_SAMPLES, seed=RANDOM_SEED)

    # 准备数据加载器
    print("Preparing data loaders...")
    train_loader, val_loader = prepare_data_loaders(train_data, val_data)

    # 创建增强世界模型
    print("Creating enhanced world model with GNN architecture...")
    model = EnhancedWorldModel(node_features=total_node_features, dropout_rate=DROPOUT_RATE)

    # 训练模型
    print("Training model with advanced strategies...")
    print(f"训练最多 {NUM_EPOCHS} 轮，早停耐心值={PATIENCE}")
    print(f"环境损失权重: {ENV_LOSS_WEIGHT}, 智能体损失权重: {AGENT_LOSS_WEIGHT}")
    print(f"环境学习率: {LEARNING_RATE}, 智能体学习率: {AGENT_LEARNING_RATE}")
    print(f"GNN层数: {GNN_NUM_LAYERS}, 隐藏通道数: {GNN_HIDDEN_CHANNELS}")
    print(f"智能体分支隐藏层: {AGENT_BRANCH_HIDDEN_SIZES}")

    trained_model, metrics = train(
        model,
        train_loader,
        val_loader,
        epochs=NUM_EPOCHS,
        env_lr=LEARNING_RATE,
        agent_lr=AGENT_LEARNING_RATE,
        weight_decay=WEIGHT_DECAY
    )

    # 保存最终模型
    print("Saving final model...")
    torch.save(trained_model.state_dict(), "models/enhanced_world_model_final.pth")

    # 评估模型
    print("Evaluating model...")
    env_acc, static_acc, agent_pos_acc, agent_mov_acc = evaluate(trained_model, test_data)

    # 打印最终测试准确率
    print("======= 最终测试结果 =======")
    print(f"环境状态准确率: {env_acc:.4f}")
    print(f"静态元素准确率: {static_acc:.4f}")
    print(f"智能体位置准确率: {agent_pos_acc:.4f}")
    print(f"智能体移动准确率: {agent_mov_acc:.4f}")
    print("==========================")

    # 运行测试episodes
    print("Running test episodes...")
    run_test_episodes(trained_model)


if __name__ == "__main__":
    main()